package Tree;

import java.util.ArrayList;
import java.util.List;

public class _863_AllNodesDistanceKinBinaryTree {
    /*private List<Integer> res = new ArrayList<>();
    private int target;
    private int K;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        this.target = target.val;
        this.K = K;
        InorderTraversal(root,root.val,1);
        return res;
    }

    public void InorderTraversal(TreeNode root,int val,int count) {
        if (root != null && count>K) {
            if (root.val == target &&count==K) {
                res.add(val);
                return;
            } else {
                InorderTraversal(root.left, val, count + 1);
                InorderTraversal(root.right, val, count + 1);
                if (root.left != null) {
                    InorderTraversal(root.left,root.left.val,1);
                }
                if (root.right != null) {
                    InorderTraversal(root.right, root.right.val, 1);
                }
            }
        }
    }*/

    List<Integer> res = new ArrayList<>();
    int K;
    int target;

    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        this.K = K;
        this.target = target.val;
        dfs(root);
        return res;
    }

    public int dfs(TreeNode root) {
        if (root == null) {
            return -1;
        } else if (root.val == target) {
            addSubTree(root, 0);
            return 1;
        } else {
            int L = dfs(root.left), R = dfs(root.right);
            if (L != -1) {
                if(L == K) res.add(root.val);
                addSubTree(root.right, L + 1);
                return L + 1;
            } else if (R != -1) {
                if (R == K) res.add(root.val);
                addSubTree(root.left, R + 1);
                return R + 1;
            } else {
                return -1;
            }
        }
    }
    public void addSubTree(TreeNode root, int dist) {
        if(root==null) return;
        if (dist == K) {
            res.add(root.val);
        } else {
            addSubTree(root.left, dist + 1);
            addSubTree(root.right, dist + 1);
        }

    }
}
